import os
import random
import string

import cv2
import easyocr  # 新增导入
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from PIL import Image, ImageDraw, ImageFont
from torch.utils.data import DataLoader, Dataset
from tqdm import trange

# 配置参数
CAPTCHA_LENGTH = 6  # 验证码字符长度
CHAR_SET = string.ascii_letters + string.digits  # 字符集（52字母+10数字）
CHAR_NUM = len(CHAR_SET)  # 字符类别数
IMAGE_SIZE = (60, 180)  # 验证码尺寸 (height, width)
DATASET_SIZE = 2000  # 总数据集大小（训练1600+测试400）
MODEL_PATH = "captcha_model.pth"  # 模型保存路径


# -------------------- 验证码生成模块 --------------------


def generate_captcha(
    text_length=CAPTCHA_LENGTH, size=(180, 60), font_size=28
):  # 适配项目参数
    """Generate simple captcha image with text and noise（来自image.py的优化版本）"""
    # 生成随机字符（使用项目统一字符集）
    captcha_text = "".join(
        random.choices(CHAR_SET, k=text_length)
    )  # 替换为项目CHAR_SET

    # 创建白底背景图像（提升OCR识别率）
    img = Image.new("RGB", size, (255, 255, 255))
    draw = ImageDraw.Draw(img)

    # 加载字体（使用系统默认，与原逻辑一致）
    try:
        font = ImageFont.truetype("arial.ttf", font_size)
    except IOError:
        font = ImageFont.load_default()

    # 绘制黑色文本（居中显示，无随机偏移保证清晰度）
    bbox = draw.textbbox((0, 0), captcha_text, font=font)
    text_width, text_height = bbox[2] - bbox[0], bbox[3] - bbox[1]
    x = (size[0] - text_width) / 2  # 移除随机偏移，保证文本居中
    y = (size[1] - text_height) / 2
    draw.text((x, y), captcha_text, fill=(0, 0, 0), font=font)  # 黑底白字→白底黑字

    # 添加少量干扰线（保留但控制数量，原image.py为3条）
    for _ in range(3):
        x1 = random.randint(0, size[0])
        y1 = random.randint(0, size[1])
        x2 = random.randint(0, size[0])
        y2 = random.randint(0, size[1])
        draw.line([(x1, y1), (x2, y2)], fill=(128, 128, 128), width=1)  # 浅灰色干扰线

    # 转换为灰度图并归一化（与原gan.py数据格式保持一致）
    img_np = np.array(img)
    img_gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    img_normalized = img_gray / 255.0
    return img_normalized, captcha_text


def generate_dataset(root_dir="dataset"):
    """生成训练/测试数据集（自动创建目录）"""
    os.makedirs(f"{root_dir}/train", exist_ok=True)
    os.makedirs(f"{root_dir}/test", exist_ok=True)

    # 生成训练集（1600张）
    for i in trange(1600, desc="Generating train set"):
        img, text = generate_captcha()
        np.save(f"{root_dir}/train/{i}.npy", img)
        with open(f"{root_dir}/train/labels.txt", "a") as f:
            f.write(f"{i}.npy,{text}\n")

    # 生成测试集（400张）
    for i in trange(400, desc="Generating test set"):
        img, text = generate_captcha()
        np.save(f"{root_dir}/test/{i}.npy", img)
        with open(f"{root_dir}/test/labels.txt", "a") as f:
            f.write(f"{i}.npy,{text}\n")


# -------------------- 数据加载模块 --------------------
class CaptchaDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.labels = {}
        with open(f"{root_dir}/labels.txt", "r") as f:
            for line in f:
                img_name, text = line.strip().split(",")
                self.labels[img_name] = text

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        img_name = list(self.labels.keys())[idx]
        img_path = os.path.join(self.root_dir, img_name)
        img = np.load(img_path)  # 加载灰度图（60, 180）

        # 新增数据增强：随机亮度调整（±10%）
        brightness_factor = 0.9 + random.random() * 0.2
        img = np.clip(img * brightness_factor, 0, 1)

        img = torch.tensor(img, dtype=torch.float32).unsqueeze(0)  # 转换为(1, 60, 180)

        text = self.labels[img_name]
        label = torch.tensor([CHAR_SET.index(c) for c in text], dtype=torch.long)
        return img, label


# -------------------- 模型定义模块 --------------------


class CaptchaModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),  # 新增BatchNorm层
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),  # 新增BatchNorm层
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),  # 新增BatchNorm层
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
        )

        self.fc = nn.Linear(128 * 7 * 22, CAPTCHA_LENGTH * CHAR_NUM)

    def forward(self, x):
        x = self.cnn(x)
        x = x.view(x.size(0), -1)  # 展平为(批量, 128*7*22)
        x = self.fc(x)
        return x.view(x.size(0), CAPTCHA_LENGTH, CHAR_NUM)  # 输出(批量, 6, 62)


# -------------------- 训练与预测模块 --------------------
def train_model():
    # 加载数据集
    train_dataset = CaptchaDataset("dataset/train")
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

    # 初始化模型
    model = CaptchaModel()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 训练循环
    for epoch in trange(20, desc="Training"):  # 简易训练20轮
        model.train()
        total_loss = 0.0
        for imgs, labels in train_loader:
            optimizer.zero_grad()
            outputs = model(imgs)  # (32,6,62)
            loss = criterion(outputs.permute(0, 2, 1), labels)  # 调整维度适应交叉熵
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

    # 保存模型
    torch.save(model.state_dict(), MODEL_PATH)


def predict_captcha(model, image):
    """使用模型+EasyOCR联合预测验证码"""
    model.eval()
    with torch.no_grad():
        # 1. 先尝试用EasyOCR识别
        reader = easyocr.Reader(["en"], gpu=False)
        ocr_result = reader.readtext(image, detail=0)
        ocr_text = "".join(ocr_result).replace(" ", "")
        if len(ocr_text) == CAPTCHA_LENGTH and all(c in CHAR_SET for c in ocr_text):
            return ocr_text  # EasyOCR有效结果直接返回

        # 2. EasyOCR失败时使用自定义模型
        img_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        img_normalized = img_gray / 255.0
        img_tensor = (
            torch.tensor(img_normalized, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        )

        output = model(img_tensor)
        _, predicted = torch.max(output, dim=2)
        predicted_text = "".join([CHAR_SET[idx] for idx in predicted[0].tolist()])
    return predicted_text


# -------------------- 主流程 --------------------
def main():
    # 生成数据集（如果不存在）
    if not os.path.exists("dataset/train/labels.txt"):
        print("Generating dataset...")
        generate_dataset()

    # 训练模型（如果未训练）
    if not os.path.exists(MODEL_PATH):
        print("Training model...")
        train_model()

    # 加载模型
    model = CaptchaModel()
    model.load_state_dict(torch.load(MODEL_PATH))

    # 生成并测试新验证码
    test_img_np, true_text = generate_captcha()
    test_img_rgb = cv2.cvtColor(
        (test_img_np * 255).astype(np.uint8), cv2.COLOR_GRAY2RGB
    )  # 转换为RGB用于显示
    recognized_text = predict_captcha(model, test_img_rgb)

    # 可视化对比
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.imshow(test_img_rgb)
    plt.title(f"GAN-Generated\nTrue: {true_text}")

    plt.subplot(1, 2, 2)
    plt.imshow(test_img_rgb)
    plt.title(f"Model Prediction\nRecognized: {recognized_text}")

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()
